import datetime
import pandas as pd
import numpy as np
from jili.core import load,save
import os
from jili.tool.state import config
from jili.data.db import getdb_client
import pandas as pd
freq={
    "k1d":"%Y",
    "k1M":"%Y",
    "k1w":"%Y",
    "k5m":"%Y%m%d",
    "k15m":"%Y%m%d",
    "k30m":"%Y%m",
    "k60m":"%Y%m",
}
rr=load(r"F:\qjs\data\right_factor.pkl")
right_feild=["open","close","high","low"]
def holdfield(k,feild):
    d=[]
    for i in k.columns:
        if i not in feild:
            d.append(i)
    if d:
        k.drop(d,axis=1)
def bingji(a,b):
    rst=[]
    for i in a:
        if i in b:
            rst.append(i)
    return rst
def df_round(k,feild,n):
    f={}
    for i in feild:
        f[i]=n
    return k.round(f)
def refeild(feild):
    for i in ["tiemkey","obj"]:
        if i not in feild:
            feild.append(i)
    return feild
def adjust(objs,k,flag,startdate=None,enddate=None):
    feild_r=bingji(right_feild,k.columns)
    if flag=="qfq":
        flags="qfq_factor_current"
    else:
        flags="hfq_factor_current"
    if isinstance(objs,str):
        if objs=="all":
            for obj in rr.keys():
                r = rr[obj]
                if startdate:
                    if enddate:
                        r = r.loc[(r.index >= startdate) & (r.index <= enddate)]
                    else:
                        r = r.loc[r.index >= startdate]
                else:
                    if enddate:
                        r = r.loc[r.index <= enddate]
                if len(r) >= 2:
                    r = r.copy(deep=True)
                    if flag != "qfq":  # 前复权，date倒叙
                        r = r.sort_index()
                        r.loc[r.index[0], flags] = 1
                        r["factor"] = np.cumprod(r[flags])
                        n = len(r.index)
                        for i in range(n):
                            if i > 0 and i < n - 1:
                                t1 = r.index[i]
                                t2 = r.index[i + 1]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1) & (k["timekey"] < t2)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                            elif i == n - 1:
                                t1 = r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                    else:
                        r.loc[r.index[0], flags] = 1
                        r["factor"] = np.cumprod(r[flags])
                        n = len(r.index)
                        for i in range(n):
                            if i > 0 and i < n - 1:
                                t1 = r.index[i]
                                t2 = r.index[i + 1]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1) & (k["timekey"] > t2)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                            elif i == n - 1:
                                t1 = r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
        else:
            obj=objs
            if obj in rr.keys():
                r = rr[obj]
                if startdate:
                    if enddate:
                        r = r.loc[(r.index >= startdate) & (r.index <= enddate)]
                    else:
                        r = r.loc[r.index >= startdate]
                else:
                    if enddate:
                        r = r.loc[r.index <= enddate]
                if len(r) >= 2:
                    r = r.copy(deep=True)
                    if flag != "qfq":  # 前复权，date倒叙
                        r = r.sort_index()
                        r.loc[r.index[0], flags] = 1
                        r["factor"] = np.cumprod(r[flags])
                        n = len(r.index)
                        for i in range(n):
                            if i > 0 and i < n - 1:
                                t1 = r.index[i]
                                t2 = r.index[i + 1]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1) & (k["timekey"] < t2)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                            elif i == n - 1:
                                t1 = r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                    else:
                        r.loc[r.index[0], flags] = 1
                        r["factor"] = np.cumprod(r[flags])
                        n = len(r.index)
                        for i in range(n):
                            if i > 0 and i < n - 1:
                                t1 = r.index[i]
                                t2 = r.index[i + 1]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1) & (k["timekey"] > t2)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                            elif i == n - 1:
                                t1 = r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
    else:
        for obj in objs:
            if obj in rr.keys():
                r=rr[obj]
                if startdate:
                    if enddate:
                        r = r.loc[(r.index >= startdate) & (r.index <= enddate)]
                    else:
                        r = r.loc[r.index >= startdate]
                else:
                    if enddate:
                        r = r.loc[r.index <= enddate]
                if len(r) >= 2:
                    r=r.copy(deep=True)
                    if flag != "qfq":#前复权，date倒叙
                        r=r.sort_index()
                        r.loc[r.index[0], flags] = 1
                        r["factor"]=np.cumprod(r[flags])
                        n=len(r.index)
                        for i in range(n):
                            if i>0 and i<n-1:
                                t1=r.index[i]
                                t2=r.index[i+1]
                                rdata=r.loc[t1,"factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1) & (k["timekey"] < t2)
                                k.loc[tt, feild_r] = k.loc[tt,feild_r] * rdata
                            elif i==n-1:
                                t1=r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] >= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                    else:
                        r.loc[r.index[0], flags] = 1
                        r["factor"] = np.cumprod(r[flags])
                        n = len(r.index)
                        for i in range(n):
                            if i > 0 and i < n - 1:
                                t1 = r.index[i]
                                t2 = r.index[i + 1]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1) & (k["timekey"] > t2)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
                            elif i == n - 1:
                                t1 = r.index[i]
                                rdata = r.loc[t1, "factor"]
                                tt = (k["obj"] == obj) & (k["timekey"] <= t1)
                                k.loc[tt, feild_r] = k.loc[tt, feild_r] * rdata
    k=df_round(k,feild_r,2)
    return k
def getstock_kdata(objs,ktype="k1d",startdate=None,enddate=None,adjustflag="hfq",feild=None,dtype="df",ip=None):
    """

    Args:
        objs: str,一个代码或是  all or list:[] 交易所代码
        ktype: str :k1d,k1dw,k1M,K5m,k15m,k30m,k60m
        startdate:str :YYYYMMDD,None:全部
        enddate:str :YYYYMMDD,None:全部
        adjustflag:str :qfq,hfq,bfq
        feild:list or none
        dtype:str:df or other

    Returns:
        objs str:单重索引的数据框
        objs list :多重索引数据框
    """
    if ip is None:
        ip=config.mongodb_ip
    db=getdb_client("kline",ip)
    rst=pd.DataFrame()
    if isinstance(objs,str):
        if objs=="all":
            k = []
            if startdate:
                if enddate:
                    starts=datetime.datetime.strptime(startdate,"%Y%m%d")
                    ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"timekey":{'$gte':starts,'$lte':ends}},refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"timekey":{'$gte':starts,'$lte':ends}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst=pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True,drop=False)
                        rst.sort_index(inplace=True)
                else:
                    starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    #ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"timekey": {'$gte': starts}}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"timekey": {'$gte': starts}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True,drop=False)
                        rst.sort_index(inplace=True)
            else:
                if enddate:
                    #starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"timekey": {'$lte': ends}}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"timekey": {'$lte': ends}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True,drop=False)
                        rst.sort_index(inplace=True)
                else:
                    #starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    #ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True,drop=False)
                        rst.sort_index(inplace=True)

        else:
            k = []
            if startdate:
                if enddate:
                    starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"obj":objs,"timekey": {'$gte': starts, '$lte': ends}}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"obj":objs,"timekey": {'$gte': starts, '$lte': ends}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                        rst.sort_index(inplace=True)
                else:
                    starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"obj":objs,"timekey": {'$gte': starts}}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"obj":objs,"timekey": {'$gte': starts}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                        rst.sort_index(inplace=True)
            else:
                if enddate:
                    # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"obj":objs,"timekey": {'$lte': ends}}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"obj":objs,"timekey": {'$lte': ends}}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                        rst.sort_index(inplace=True)
                else:
                    # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                    # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                    cu_name = "stock_" + ktype
                    cu = db[cu_name]
                    if feild:
                        for i in cu.find({"obj":objs}, refeild(feild)).batch_size(1000):
                            k.append(i)
                    else:
                        for i in cu.find({"obj":objs}).batch_size(1000):
                            k.append(i)
                    if k:
                        rst = pd.DataFrame(k)
                        rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                        rst.sort_index(inplace=True)
    elif isinstance(objs,list):
        k = []
        if startdate:
            if enddate:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "stock_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj":{"$in":objs} , "timekey": {'$gte': starts, '$lte': ends}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj":{"$in":objs} , "timekey": {'$gte': starts, '$lte': ends}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "stock_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in":objs}, "timekey": {'$gte': starts}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in":objs}, "timekey": {'$gte': starts}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
        else:
            if enddate:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "stock_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in":objs}, "timekey": {'$lte': ends}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in":objs}, "timekey": {'$lte': ends}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "stock_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in":objs}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in":objs}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
    db.client.close()
    if adjustflag!="bfq":
        rst=adjust(objs,rst,adjustflag,startdate,enddate)
    if dtype != "df":
        rst = rst.to_dict("index")
    return rst

def getindex_kdata(objs,ktype="k1d",startdate=None,enddate=None,feild=None,dtype="df",ip=None):
    """

    Args:
        objs: str,一个代码或是  or list:[] 交易所代码
        ktype: str :k1d,k1dw,k1M
        startdate:str :YYYYMMDD,None:全部
        enddate:str :YYYYMMDD,None:全部
        feild:list or none
        dtype:str:df or other

    Returns:
        objs str:单重索引的数据框
        objs list :多重索引数据框
    """
    if ip   is None:
        ip = config.mongodb_ip
    db = getdb_client("kline", ip)
    rst=pd.DataFrame()
    if isinstance(objs,str):
        k = []
        if startdate:
            if enddate:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": objs, "timekey": {'$gte': starts, '$lte': ends}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": objs, "timekey": {'$gte': starts, '$lte': ends}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": objs, "timekey": {'$gte': starts}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": objs, "timekey": {'$gte': starts}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
        else:
            if enddate:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": objs, "timekey": {'$lte': ends}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": objs, "timekey": {'$lte': ends}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": objs}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": objs}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
    elif isinstance(objs,list):
        k = []
        if startdate:
            if enddate:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$gte': starts, '$lte': ends}},
                                     refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$gte': starts, '$lte': ends}}).sort(
                        [("timekey", 1)]).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$gte': starts}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$gte': starts}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
        else:
            if enddate:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$lte': ends}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in": objs}, "timekey": {'$lte': ends}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
            else:
                # starts = datetime.datetime.strptime(startdate, "%Y%m%d")
                # ends = datetime.datetime.strptime(enddate, "%Y%m%d")
                cu_name = "index_" + ktype
                cu = db[cu_name]
                if feild:
                    for i in cu.find({"obj": {"$in": objs}}, refeild(feild)).batch_size(1000):
                        k.append(i)
                else:
                    for i in cu.find({"obj": {"$in": objs}}).batch_size(1000):
                        k.append(i)
                if k:
                    rst = pd.DataFrame(k)
                    rst.set_index(["timekey", "obj"], inplace=True, drop=False)
                    rst.sort_index(inplace=True)
    db.client.close()
    if dtype != "df":
        rst = rst.to_dict("index")
    return rst